#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PoseGuard GUI
左侧：人脸库管理（上传图片、从摄像头抓拍、输入姓名、训练模型、列表显示）
右侧：摄像头预览（集成姿态/行为与人脸识别，识别到则用姓名代替ID）
"""

import os
import threading
import time
from typing import Optional

import cv2
import tkinter as tk
from tkinter import ttk, filedialog, messagebox, simpledialog
from PIL import Image, ImageTk

from main import PoseAssessmentApp
from face_recognizer import FaceRecognizer


class PoseGuardGUI(tk.Tk):
    def __init__(self, camera_id: int = 0):
        super().__init__()
        self.title("PoseGuard GUI - 人脸识别 + 姿态评估")
        self.geometry("1100x650")
        self.resizable(True, True)

        # 左侧面板（人脸库管理）
        self.left_frame = ttk.Frame(self, width=360)
        self.left_frame.pack(side=tk.LEFT, fill=tk.Y)

        # 右侧面板（摄像头预览）
        self.right_frame = ttk.Frame(self)
        self.right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

        # 人脸识别与应用
        self.recognizer = FaceRecognizer()
        self.app = PoseAssessmentApp()
        # 将同一个识别器实例注入主程序，保证训练/识别一致
        try:
            self.app.face_recognizer = self.recognizer
        except Exception:
            pass

        # 摄像头
        self.camera_id = camera_id
        self.cap: Optional[cv2.VideoCapture] = None
        self.preview_label = ttk.Label(self.right_frame)
        self.preview_label.pack(fill=tk.BOTH, expand=True)
        self.running = False
        self.preview_thread: Optional[threading.Thread] = None
        self._ui_job = None
        self._last_frame_pil: Optional[Image.Image] = None

        self._build_left_ui()

        self.protocol("WM_DELETE_WINDOW", self.on_close)

    def _show_status_message(self, text: str):
        try:
            # 清除图片，显示文本提示
            self.preview_label.configure(image="", text=text, anchor=tk.CENTER, justify=tk.CENTER)
            self.preview_label.imgtk = None
        except Exception:
            pass

    def _build_left_ui(self):
        # 操作按钮
        frm_btns = ttk.Frame(self.left_frame)
        frm_btns.pack(fill=tk.X, padx=10, pady=(12, 6))

        ttk.Button(frm_btns, text="上传图片并加入库", command=self.on_add_from_file).pack(fill=tk.X, pady=3)
        ttk.Button(frm_btns, text="从摄像头抓拍加入库", command=self.on_add_from_camera).pack(fill=tk.X, pady=3)
        ttk.Button(frm_btns, text="训练模型", command=self.on_train).pack(fill=tk.X, pady=3)
        ttk.Button(frm_btns, text="移除所选人员", command=self.on_remove_selected).pack(fill=tk.X, pady=3)
        ttk.Button(frm_btns, text="刷新列表", command=self.refresh_people_list).pack(fill=tk.X, pady=3)
        ttk.Separator(self.left_frame, orient=tk.HORIZONTAL).pack(fill=tk.X, padx=10, pady=8)

        # 人员列表
        ttk.Label(self.left_frame, text="已有人脸条目:").pack(anchor=tk.W, padx=10)
        self.people_list = tk.Listbox(self.left_frame, height=18)
        self.people_list.pack(fill=tk.BOTH, expand=False, padx=10, pady=(4, 8))

        # 预览控制
        ttk.Separator(self.left_frame, orient=tk.HORIZONTAL).pack(fill=tk.X, padx=10, pady=8)
        frm_cam = ttk.Frame(self.left_frame)
        frm_cam.pack(fill=tk.X, padx=10, pady=6)
        ttk.Button(frm_cam, text="启动预览", command=self.start_preview).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))
        ttk.Button(frm_cam, text="停止预览", command=self.stop_preview).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(5, 0))

        self.refresh_people_list()

    def on_add_from_file(self):
        paths = filedialog.askopenfilenames(title="选择人脸图片", filetypes=[("图片", "*.png;*.jpg;*.jpeg")])
        if not paths:
            return
        # 选择后再询问姓名（支持多人的方式：多次上传分别输入）
        name = simpledialog.askstring("输入姓名", "请输入该批图片对应的姓名：", parent=self)
        if not name or not name.strip():
            messagebox.showwarning("提示", "未输入姓名，已取消")
            return
        name = name.strip()
        added = 0
        for p in paths:
            try:
                img = cv2.imread(p)
                if img is None:
                    continue
                if self.recognizer.add_face_image(name, img):
                    added += 1
            except Exception:
                pass
        self.refresh_people_list()
        messagebox.showinfo("完成", f"已添加 {added} 张样本到 {name}")
        # 重建 pHash 索引，支持免训练直匹配
        try:
            self.recognizer.rebuild_index()
        except Exception:
            pass

    def on_add_from_camera(self):
        # 使用当前帧抓拍
        frame = getattr(self, "_last_frame_bgr", None)
        if frame is None:
            messagebox.showwarning("提示", "当前没有可用的摄像头画面，请先启动预览")
            return
        name = simpledialog.askstring("输入姓名", "请输入该抓拍对应的姓名：", parent=self)
        if not name or not name.strip():
            messagebox.showwarning("提示", "未输入姓名，已取消")
            return
        name = name.strip()
        if self.recognizer.add_face_image(name, frame) is None:
            messagebox.showwarning("提示", "未能检测到人脸，建议换一张角度/光线更好的图像")
            return
        self.refresh_people_list()
        messagebox.showinfo("完成", f"已从摄像头抓拍添加 1 张到 {name}")
        try:
            self.recognizer.rebuild_index()
        except Exception:
            pass

    def on_train(self):
        ok = self.recognizer.train()
        if ok:
            messagebox.showinfo("训练完成", "人脸识别模型已训练并保存（建议每人≥10张样本以提升稳定性）")
            # 训练完成后清空绑定
            try:
                self.app.clear_bindings()
            except Exception:
                pass
        else:
            messagebox.showwarning("训练失败", "请先添加样本（至少1人且每人≥1张）并确保已安装 opencv-contrib-python")

    def on_remove_selected(self):
        try:
            sel = self.people_list.curselection()
            if not sel:
                messagebox.showinfo("提示", "请先在列表中选择一名人员")
                return
            name = self.people_list.get(sel[0])
            if not name:
                return
            if messagebox.askyesno("确认", f"确定要删除人员及其样本：{name}？此操作不可撤销"):
                # 删除 faces/<name> 目录、更新 labels.json 与 inverse 映射，并删除模型文件（迫使重新训练）
                base_dir = self.recognizer.base_dir
                person_dir = os.path.join(base_dir, name)
                try:
                    # 移除目录
                    if os.path.isdir(person_dir):
                        for fn in os.listdir(person_dir):
                            try:
                                os.remove(os.path.join(person_dir, fn))
                            except Exception:
                                pass
                        try:
                            os.rmdir(person_dir)
                        except Exception:
                            pass
                    # 更新 labels
                    lid = self.recognizer.labels.pop(name, None)
                    if lid is not None:
                        self.recognizer.inverse_labels.pop(lid, None)
                    self.recognizer._save_labels()
                    # 删除已训练模型，避免使用旧模型
                    try:
                        if os.path.exists(self.recognizer.model_path):
                            os.remove(self.recognizer.model_path)
                    except Exception:
                        pass
                    # 重新初始化识别器，确保状态一致
                    self.recognizer._init_recognizer()
                    self.recognizer._try_load_model()
                except Exception as e:
                    messagebox.showerror("错误", f"删除失败: {e}")
                self.refresh_people_list()
                messagebox.showinfo("完成", f"已删除：{name}。请重新训练模型以生效。")
                try:
                    self.recognizer.rebuild_index()
                except Exception:
                    pass
                # 删除人员后清空绑定
                try:
                    self.app.clear_bindings()
                except Exception:
                    pass
        except Exception:
            messagebox.showerror("错误", "删除过程出现异常")

    def refresh_people_list(self):
        self.people_list.delete(0, tk.END)
        try:
            names = sorted(self.recognizer.labels.keys())
        except Exception:
            names = []
        for n in names:
            self.people_list.insert(tk.END, n)

    def start_preview(self):
        if self.running:
            return
        if self.cap is None:
            self.cap = cv2.VideoCapture(self.camera_id)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
            # 未能打开摄像头时给出提示
            try:
                if not self.cap.isOpened():
                    self._show_status_message("请检索摄像头是不是占用，暂无信息")
                    try:
                        # 释放资源
                        self.cap.release()
                    except Exception:
                        pass
                    self.cap = None
                    return
            except Exception:
                self._show_status_message("请检索摄像头是不是占用，暂无信息")
                return
        self.running = True
        self.preview_thread = threading.Thread(target=self._preview_loop, daemon=True)
        self.preview_thread.start()
        # 启动主线程UI定时刷新，降低闪烁
        self._start_ui_timer()

    def stop_preview(self):
        self.running = False
        # 取消UI定时器
        try:
            if self._ui_job is not None:
                self.after_cancel(self._ui_job)
        except Exception:
            pass
        self._ui_job = None

    def _preview_loop(self):
        read_fail_count = 0
        while self.running:
            try:
                if self.cap is None:
                    break
                ret, frame = self.cap.read()
                if not ret:
                    read_fail_count += 1
                    if read_fail_count >= 15:
                        # 连续多次读取失败，显示提示
                        self._show_status_message("请检索摄像头是不是占用，暂无信息")
                        time.sleep(0.2)
                    else:
                        time.sleep(0.02)
                    continue
                # 成功读取后重置失败计数
                read_fail_count = 0
                # 保存 BGR 帧供抓拍使用
                self._last_frame_bgr = frame.copy()
                # 处理图像（后台线程仅准备 PIL 图像，不直接操作 Tk）
                result_image, _ = self.app.process_image(frame)
                rgb = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
                self._last_frame_pil = Image.fromarray(rgb)
            except Exception:
                time.sleep(0.02)
                continue
        # 清理
        try:
            self.preview_label.configure(image="")
            self.preview_label.imgtk = None
        except Exception:
            pass

    def _start_ui_timer(self):
        # 主线程以固定频率（~30fps）刷新，减少闪烁
        try:
            self._ui_job = self.after(33, self._ui_update_tick)
        except Exception:
            self._ui_job = None

    def _ui_update_tick(self):
        try:
            if self._last_frame_pil is not None:
                imgtk = ImageTk.PhotoImage(image=self._last_frame_pil)
                self.preview_label.imgtk = imgtk
                self.preview_label.configure(image=imgtk, text="")
        except Exception:
            pass
        finally:
            if self.running:
                try:
                    self._ui_job = self.after(33, self._ui_update_tick)
                except Exception:
                    self._ui_job = None

    def on_close(self):
        self.stop_preview()
        try:
            if self.cap is not None:
                self.cap.release()
        except Exception:
            pass
        # 关闭时也给出空闲提示
        try:
            self._show_status_message("请检索摄像头是不是占用，暂无信息")
        except Exception:
            pass
        self.destroy()


def main():
    app = PoseGuardGUI(camera_id=0)
    app.mainloop()


if __name__ == "__main__":
    main()


